#https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=9145522&tag=1 的LSTM
class LSTM(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim, num_layers, dropout_rate):
        super(LSTM, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.lstm1 = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers, batch_first=True)
        self.dropout1 = nn.Dropout(dropout_rate)
       
        self.lstm2 = nn.LSTM(hidden_dim, hidden_dim, num_layers=num_layers, batch_first=True)
        self.dropout2 = nn.Dropout(dropout_rate)
       
        self.fc = nn.Linear(hidden_dim, output_dim)
        
    def forward(self, x):
        #print('\ninput:',x.shape) #([128, 10, 5])
        #x = x.permute(0, 2, 1)  ???
        x,_ = self.lstm1(x)
        x = self.dropout1(x)
        x,_ = self.lstm2(x)
        x = self.dropout2(x)
        x = self.fc(x[:,-1,:]) #全连接层 #这里x[:,-1,:]没太懂
        return x

#https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=9145522&tag=1 的CNNLSTM(应该
class CNNLSTM(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim=50, dropout_rate=0.3):
        super(CNNLSTM, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=input_dim, out_channels=32, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool1d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv1d(in_channels=32, out_channels=32, kernel_size=3,padding=1)
       
        self.lstm1 = nn.LSTM(input_size=32, hidden_size=50)
        self.lstm2 = nn.LSTM(input_size=50, hidden_size=50 )
        self.dropout = nn.Dropout(dropout_rate)
        self.fc = nn.Linear(50,output_dim) #本来是(50,4)->(50,3)

    def forward(self, x):
        x = x.permute(0, 2, 1) 
        
        x = self.relu(self.conv1(x))
        x = self.pool(x)
        x = self.relu(self.conv2(x))
        x = self.pool(x)

        x = x.permute(0, 2, 1)  
        x,_ = self.lstm1(x)
        x = self.dropout(x)
        x,_ = self.lstm2(x) 
        x = x[:, -1, :]
        x = self.dropout(x)

        x = self.fc(x) #全连接层
        return x